#!/usr/bin/env python3
"""
模型结构检查脚本
用法: python scripts/inspect_model.py
"""
import sys
from pathlib import Path

sys.path.append(str(Path(__file__).parent.parent))

import torch
from src.model import JiaboForCausalLM, JiaboModelConfig
from src.tokenizer import JiaboTokenizer
from src.utils import load_json_config


def main():
    print("🔍 Jiabo-0.5B-R1 模型结构检查")
    
    # 加载模型
    model_cfg = load_json_config("configs/model_config.json")
    config = JiaboModelConfig(**model_cfg)
    model = JiaboForCausalLM(config)
    
    # 打印模型架构
    print("\n📐 模型配置:")
    for key, value in model_cfg.items():
        print(f"  {key}: {value}")
    
    print("\n🧩 模型层数:")
    print(model)
    
    # 参数量统计
    total_params = sum(p.numel() for p in model.parameters())
    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
    
    print(f"\n📊 参数量统计:")
    print(f"  总参数量: {total_params:,} ({total_params / 1e6:.2f}M)")
    print(f"  可训练参数量: {trainable_params:,} ({trainable_params / 1e6:.2f}M)")
    
    # 测试前向传播
    tokenizer = JiaboTokenizer("data/vocab.json")
    dummy_input = torch.tensor([tokenizer.encode("Hello world")])
    
    print(f"\n🚀 测试前向传播:")
    print(f"  输入shape: {dummy_input.shape}")
    
    output, loss = model(dummy_input, labels=dummy_input)
    print(f"  输出logits shape: {output.shape}")
    print(f"  示例loss: {loss.item():.4f}")


if __name__ == "__main__":
    main()
